import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix, precision_recall_fscore_support, accuracy_score
from sklearn.preprocessing import StandardScaler
import joblib
from joblib import dump
import os

# 加载数据
data = pd.read_csv('Normalization_total_extend.csv')
X = data.iloc[:, :-1]  # 假设最后一列是类别标签
y = data.iloc[:, -1]

# 划分数据集为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 特征缩放（尽管已经标准化，但再次确认总是好的）
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 定义逻辑回归模型
logreg = LogisticRegression(multi_class='multinomial', solver='lbfgs', max_iter=1000, random_state=42)

# 检查模型文件是否存在
model_file = 'LogisticRegressionModel/logistic_regression_model.pkl'
if os.path.exists(model_file):
    print("Model file exists, loading model...")
    logreg = joblib.load(model_file)
else:
    print("Model file does not exist, training model...")
    # 训练模型
    logreg.fit(X_train_scaled, y_train)

    # 保存模型
    joblib.dump(logreg, model_file)
    print("Model trained and saved.")

# 模型评估
y_pred = logreg.predict(X_test_scaled)
accuracy = accuracy_score(y_test, y_pred)
conf_mat = confusion_matrix(y_test, y_pred)
precision, recall, fscore, _ = precision_recall_fscore_support(y_test, y_pred, average='weighted')

print("Accuracy:", accuracy)
print("Confusion Matrix:\n", conf_mat)
print("Precision:", precision)
print("Recall:", recall)
print("F1-score:", fscore)